import os
import csv
import numpy as np
from PIL import Image

import matplotlib.pyplot as plt
import copy

def evaluteTop1_5(classfication, lines, metrics_out_path):
    correct_1 = 0
    correct_5 = 0
    #   创建类别数目乘类别数目的混淆矩阵
    conf_mat = np.zeros([classfication.num_classes, classfication.num_classes], dtype=int)
    #   创建 处理每张图片的所花费的时间的 列表
    time_list = []

    total = len(lines)
    for index, line in enumerate(lines):

        annotation_data = line.rstrip().split()
        x = Image.open(annotation_data[0])
        y = int(annotation_data[1])

        pred, image_time = classfication.detect_image(x)
        time_list.append(image_time)

        pred_1      = np.argmax(pred)
        correct_1   += pred_1 == y

        pred_5      = np.argsort(pred)[::-1]
        pred_5      = pred_5[:5]
        correct_5   += y in pred_5

        conf_mat[y, pred_1] += 1

        if index % 100 == 0:
            print("[%d/%d] " % (index, total))

    # 删除模型处理第一张图片花费的时间
    del time_list[0]
    time_array = np.array(time_list)
    fps = 1 / np.mean(time_array)

    Recall      = per_class_Recall(conf_mat)
    Precision   = per_class_Precision(conf_mat)

    show_results(metrics_out_path, conf_mat, Recall, Precision, classfication.class_names)
    show_confmat(conf_mat, classfication.class_names, metrics_out_path)

    return correct_1 / total, correct_5 / total, Recall, Precision, fps

def per_class_Recall(hist):
    return np.diag(hist) / np.maximum(hist.sum(1), 1)

def per_class_Precision(hist):
    return np.diag(hist) / np.maximum(hist.sum(0), 1)

def adjust_axes(r, t, fig, axes):
    bb                  = t.get_window_extent(renderer=r)
    text_width_inches   = bb.width / fig.dpi
    current_fig_width   = fig.get_figwidth()
    new_fig_width       = current_fig_width + text_width_inches
    propotion           = new_fig_width / current_fig_width
    x_lim               = axes.get_xlim()
    axes.set_xlim([x_lim[0], x_lim[1] * propotion])

def draw_plot_func(values, name_classes, plot_title, x_label, output_path, tick_font_size = 12, plt_show = True):
    fig     = plt.gcf()
    axes    = plt.gca()
    plt.barh(range(len(values)), values, color='royalblue')
    plt.title(plot_title, fontsize=tick_font_size + 2)
    plt.xlabel(x_label, fontsize=tick_font_size)
    plt.yticks(range(len(values)), name_classes, fontsize=tick_font_size)
    r = fig.canvas.get_renderer()
    for i, val in enumerate(values):
        str_val = " " + str(val)
        if val < 1.0:
            str_val = " {0:.2f}".format(val)
        t = plt.text(val, i, str_val, color='royalblue', va='center', fontweight='bold')
        if i == (len(values)-1):
            adjust_axes(r, t, fig, axes)

    fig.tight_layout()
    fig.savefig(output_path)
    if plt_show:
        plt.show()
    plt.close()

def show_results(miou_out_path, hist, Recall, Precision, name_classes, tick_font_size = 12):
    draw_plot_func(Recall, name_classes, "mRecall = {0:.2f}%".format(np.nanmean(Recall)*100), "Recall", \
        os.path.join(miou_out_path, "Recall.png"), tick_font_size = tick_font_size, plt_show = False)
    print("Save Recall out to " + os.path.join(miou_out_path, "Recall.png"))

    draw_plot_func(Precision, name_classes, "mPrecision = {0:.2f}%".format(np.nanmean(Precision)*100), "Precision", \
        os.path.join(miou_out_path, "Precision.png"), tick_font_size = tick_font_size, plt_show = False)
    print("Save Precision out to " + os.path.join(miou_out_path, "Precision.png"))

    with open(os.path.join(miou_out_path, "confusion_matrix.csv"), 'w', newline='') as f:
        writer          = csv.writer(f)
        writer_list     = []
        writer_list.append([' '] + [str(c) for c in name_classes])
        for i in range(len(hist)):
            writer_list.append([name_classes[i]] + [str(x) for x in hist[i]])
        writer.writerows(writer_list)
    print("Save confusion_matrix out to " + os.path.join(miou_out_path, "confusion_matrix.csv"))

def show_confmat(confusion_matrix, classes_name, out_dir, font_size = 12):
    """
    可视化混淆矩阵，保存png格式
    :param confusion_matrix: nd-array
    :param classes_name: list,各类别名称
    :param out_dir: str, png输出的文件夹
    :return:
    """

    from matplotlib import rcParams
    fontconfig = {
        "font.family": 'Times New Roman',  # 设置字体类型
    }
    rcParams.update(fontconfig)

    xiters = range(len(classes_name))

    # 归一化,转为float类型
    confusion_matrix_n = confusion_matrix.copy().astype(np.float64)
    for i in xiters:
        confusion_matrix_n[i, :] = confusion_matrix[i, :] / confusion_matrix[i, :].sum()

    # cmap = plt.cm.get_cmap('Greys')  # 获取颜色
    # 按照像素显示出矩阵
    plt.imshow(confusion_matrix_n, interpolation='nearest', cmap=plt.cm.Blues)
    # 改变颜色: ('Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds','YlOrBr', 'YlOrRd',
    # 'OrRd', 'PuRd', 'RdPu', 'BuPu','GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn')
    plt.colorbar()

    #------------#
    # 设置文字
    #------------#
    # 设置坐标刻度
    # plt.xticks(xiters, classes_name, fontsize=font_size)
    plt.xticks(xiters, classes_name, fontsize=font_size, rotation=60)  # rotation设置字的旋转角度
    plt.yticks(xiters, classes_name, fontsize=font_size)
    # 设置坐标轴标题
    plt.xlabel('Predict label', fontsize=font_size + 2)
    plt.ylabel('True label', fontsize=font_size + 2)
    plt.title('confusion matrix', fontsize=font_size + 2)

    # 打印数字
    for i in range(confusion_matrix_n.shape[0]):
        for j in range(confusion_matrix_n.shape[1]):
            # 如果是对角线位置
            if (i == j):
                plt.text(x=j, y=i, s=int(confusion_matrix[i, j]), va='center', ha='center',
                         color='white', fontsize=font_size)
            else:
                plt.text(x=j, y=i, s=int(confusion_matrix[i, j]), va='center', ha='center', fontsize=font_size)

    plt.tight_layout()
    # 保存
    plt.savefig(os.path.join(out_dir, 'confusion_matrix.png'))
    plt.close()
    print("Save confusion_matrix out to " + os.path.join(out_dir, 'confusion_matrix.png'))
